"""
WebSocket服务单元测试
==================

测试WebSocketService的所有功能，包括连接管理、消息广播和用户订阅
"""

import pytest
from unittest.mock import AsyncMock, Mock, patch, MagicMock
import asyncio
import json
from datetime import datetime, timezone
from decimal import Decimal

# Mock WebSocketService for testing
from unittest.mock import Mock

# Mock WebSocketService implementation for testing
class WebSocketService:
    """WebSocket service mock for testing"""
    def __init__(self):
        self.active_connections = {}
        self.message_log = []
    
    async def connect(self, websocket, user_id):
        """Mock connect method"""
        if user_id not in self.active_connections:
            self.active_connections[user_id] = []
        self.active_connections[user_id].append(websocket)
        return True
    
    async def disconnect(self, websocket, user_id):
        """Mock disconnect method"""
        if user_id in self.active_connections:
            try:
                self.active_connections[user_id].remove(websocket)
                if not self.active_connections[user_id]:
                    del self.active_connections[user_id]
                return True
            except ValueError:
                pass
        return False
    
    async def send_to_user(self, user_id, message):
        """Mock send to user method"""
        if user_id in self.active_connections:
            self.message_log.append((user_id, message))
            return True
        return False
    
    async def broadcast(self, message, user_ids=None):
        """Mock broadcast method"""
        if user_ids is None:
            user_ids = list(self.active_connections.keys())
        
        success_count = 0
        for user_id in user_ids:
            if user_id in self.active_connections:
                self.message_log.append((user_id, message))
                success_count += 1
        return success_count
    
    def get_active_connections(self):
        """Mock get active connections count"""
        return sum(len(connections) for connections in self.active_connections.values())
    
    def is_user_connected(self, user_id):
        """Mock is user connected check"""
        return user_id in self.active_connections and len(self.active_connections[user_id]) > 0
    
    def get_user_connection_count(self, user_id):
        """Mock get user connection count"""
        return len(self.active_connections.get(user_id, []))


@pytest.mark.unit
class TestWebSocketService:
    """WebSocket服务单元测试类"""
    
    @pytest.fixture
    def mock_websocket(self):
        """Mock WebSocket连接"""
        websocket = AsyncMock()
        websocket.send = AsyncMock()
        websocket.recv = AsyncMock()
        websocket.close = AsyncMock()
        websocket.closed = False
        websocket.client_state = Mock()
        websocket.client_state.value = 1  # CONNECTED状态
        return websocket
    
    @pytest.fixture
    def websocket_service(self):
        """创建WebSocket服务实例"""
        return WebSocketService()
    
    @pytest.mark.asyncio
    async def test_connect_new_user(self, websocket_service, mock_websocket):
        """测试新用户连接"""
        user_id = 123
        
        # 连接用户
        await websocket_service.connect(mock_websocket, user_id)
        
        # 验证连接被添加
        assert user_id in websocket_service.active_connections
        assert mock_websocket in websocket_service.active_connections[user_id]
        assert websocket_service.get_active_connections() == 1
    
    @pytest.mark.asyncio
    async def test_connect_existing_user_multiple_connections(self, websocket_service, mock_websocket):
        """测试现有用户的多重连接"""
        user_id = 123
        mock_websocket_2 = AsyncMock()
        
        # 连接同一用户的两个WebSocket
        await websocket_service.connect(mock_websocket, user_id)
        await websocket_service.connect(mock_websocket_2, user_id)
        
        # 验证两个连接都被保存
        assert user_id in websocket_service.active_connections
        assert len(websocket_service.active_connections[user_id]) == 2
        assert mock_websocket in websocket_service.active_connections[user_id]
        assert mock_websocket_2 in websocket_service.active_connections[user_id]
    
    @pytest.mark.asyncio
    async def test_disconnect_existing_user(self, websocket_service, mock_websocket):
        """测试断开现有用户连接"""
        user_id = 123
        
        # 先连接用户
        await websocket_service.connect(mock_websocket, user_id)
        assert websocket_service.get_active_connections() == 1
        
        # 断开连接
        await websocket_service.disconnect(mock_websocket, user_id)
        
        # 验证连接被移除
        assert user_id not in websocket_service.active_connections
        assert websocket_service.get_active_connections() == 0
        mock_websocket.close.assert_called_once()
    
    @pytest.mark.asyncio
    async def test_disconnect_user_with_multiple_connections(self, websocket_service, mock_websocket):
        """测试断开有多个连接的用户的单个连接"""
        user_id = 123
        mock_websocket_2 = AsyncMock()
        
        # 连接同一用户的两个WebSocket
        await websocket_service.connect(mock_websocket, user_id)
        await websocket_service.connect(mock_websocket_2, user_id)
        
        # 断开其中一个连接
        await websocket_service.disconnect(mock_websocket, user_id)
        
        # 验证只有一个连接被移除
        assert user_id in websocket_service.active_connections
        assert len(websocket_service.active_connections[user_id]) == 1
        assert mock_websocket_2 in websocket_service.active_connections[user_id]
        assert mock_websocket not in websocket_service.active_connections[user_id]
        mock_websocket.close.assert_called_once()
    
    @pytest.mark.asyncio
    async def test_disconnect_nonexistent_user(self, websocket_service, mock_websocket):
        """测试断开不存在的用户连接"""
        user_id = 999
        
        # 尝试断开不存在的连接，不应该抛出异常
        await websocket_service.disconnect(mock_websocket, user_id)
        
        # 验证状态没有变化
        assert websocket_service.get_active_connections() == 0
        mock_websocket.close.assert_called_once()
    
    @pytest.mark.asyncio
    async def test_send_to_user_existing(self, websocket_service, mock_websocket):
        """测试向存在的用户发送消息"""
        user_id = 123
        message = {"type": "market_data", "symbol": "AAPL", "price": 150.25}
        
        # 先连接用户
        await websocket_service.connect(mock_websocket, user_id)
        
        # 发送消息
        success = await websocket_service.send_to_user(user_id, message)
        
        # 验证消息发送
        assert success is True
        expected_message = json.dumps(message)
        mock_websocket.send.assert_called_once_with(expected_message)
    
    @pytest.mark.asyncio
    async def test_send_to_user_multiple_connections(self, websocket_service, mock_websocket):
        """测试向有多个连接的用户发送消息"""
        user_id = 123
        mock_websocket_2 = AsyncMock()
        message = {"type": "market_data", "symbol": "AAPL", "price": 150.25}
        
        # 连接同一用户的两个WebSocket
        await websocket_service.connect(mock_websocket, user_id)
        await websocket_service.connect(mock_websocket_2, user_id)
        
        # 发送消息
        success = await websocket_service.send_to_user(user_id, message)
        
        # 验证消息发送到所有连接
        assert success is True
        expected_message = json.dumps(message)
        mock_websocket.send.assert_called_once_with(expected_message)
        mock_websocket_2.send.assert_called_once_with(expected_message)
    
    @pytest.mark.asyncio
    async def test_send_to_user_nonexistent(self, websocket_service):
        """测试向不存在的用户发送消息"""
        user_id = 999
        message = {"type": "market_data", "symbol": "AAPL", "price": 150.25}
        
        # 发送消息给不存在的用户
        success = await websocket_service.send_to_user(user_id, message)
        
        # 验证返回False
        assert success is False
    
    @pytest.mark.asyncio
    async def test_send_to_user_connection_error(self, websocket_service, mock_websocket):
        """测试发送消息时连接错误"""
        user_id = 123
        message = {"type": "market_data", "symbol": "AAPL", "price": 150.25}
        
        # 设置WebSocket发送失败
        mock_websocket.send.side_effect = ConnectionResetError("Connection lost")
        
        # 先连接用户
        await websocket_service.connect(mock_websocket, user_id)
        
        # 发送消息
        success = await websocket_service.send_to_user(user_id, message)
        
        # 验证处理错误并移除坏连接
        assert success is False
        # 验证坏连接被移除
        assert user_id not in websocket_service.active_connections or \
               mock_websocket not in websocket_service.active_connections.get(user_id, [])
    
    @pytest.mark.asyncio
    async def test_broadcast_to_all_users(self, websocket_service):
        """测试向所有用户广播消息"""
        # 创建多个用户连接
        user_connections = {}
        for user_id in [1, 2, 3]:
            mock_ws = AsyncMock()
            user_connections[user_id] = mock_ws
            await websocket_service.connect(mock_ws, user_id)
        
        message = {"type": "market_update", "message": "Market is open"}
        
        # 广播消息
        success_count = await websocket_service.broadcast(message)
        
        # 验证所有用户收到消息
        assert success_count == 3
        expected_message = json.dumps(message)
        for mock_ws in user_connections.values():
            mock_ws.send.assert_called_once_with(expected_message)
    
    @pytest.mark.asyncio
    async def test_broadcast_to_specific_users(self, websocket_service):
        """测试向特定用户列表广播消息"""
        # 创建多个用户连接
        all_users = [1, 2, 3, 4, 5]
        target_users = [1, 3, 5]
        user_connections = {}
        
        for user_id in all_users:
            mock_ws = AsyncMock()
            user_connections[user_id] = mock_ws
            await websocket_service.connect(mock_ws, user_id)
        
        message = {"type": "targeted_update", "message": "Special announcement"}
        
        # 向特定用户广播
        success_count = await websocket_service.broadcast(message, user_ids=target_users)
        
        # 验证只有目标用户收到消息
        assert success_count == 3
        expected_message = json.dumps(message)
        
        for user_id in target_users:
            user_connections[user_id].send.assert_called_once_with(expected_message)
        
        for user_id in [2, 4]:  # 非目标用户
            user_connections[user_id].send.assert_not_called()
    
    @pytest.mark.asyncio
    async def test_broadcast_empty_connections(self, websocket_service):
        """测试在没有连接时广播消息"""
        message = {"type": "market_update", "message": "No one to receive"}
        
        # 广播消息
        success_count = await websocket_service.broadcast(message)
        
        # 验证返回0
        assert success_count == 0
    
    @pytest.mark.asyncio
    async def test_broadcast_with_connection_errors(self, websocket_service):
        """测试广播时部分连接错误"""
        # 创建用户连接，其中一些会失败
        working_users = [1, 3]
        failing_users = [2, 4]
        
        for user_id in working_users:
            mock_ws = AsyncMock()
            await websocket_service.connect(mock_ws, user_id)
        
        for user_id in failing_users:
            mock_ws = AsyncMock()
            mock_ws.send.side_effect = ConnectionResetError("Connection failed")
            await websocket_service.connect(mock_ws, user_id)
        
        message = {"type": "market_update", "message": "Test broadcast"}
        
        # 广播消息
        success_count = await websocket_service.broadcast(message)
        
        # 验证只有正常连接成功
        assert success_count == 2
        
        # 验证失败的连接被清理
        for user_id in failing_users:
            assert user_id not in websocket_service.active_connections or \
                   len(websocket_service.active_connections[user_id]) == 0
    
    def test_get_active_connections_count(self, websocket_service):
        """测试获取活跃连接数"""
        # 初始状态
        assert websocket_service.get_active_connections() == 0
        
        # 手动添加连接来测试计数
        websocket_service.active_connections[1] = [AsyncMock(), AsyncMock()]
        websocket_service.active_connections[2] = [AsyncMock()]
        websocket_service.active_connections[3] = [AsyncMock(), AsyncMock(), AsyncMock()]
        
        # 验证总连接数
        assert websocket_service.get_active_connections() == 6
    
    def test_is_user_connected_true(self, websocket_service):
        """测试用户连接状态检查（已连接）"""
        user_id = 123
        websocket_service.active_connections[user_id] = [AsyncMock()]
        
        assert websocket_service.is_user_connected(user_id) is True
    
    def test_is_user_connected_false(self, websocket_service):
        """测试用户连接状态检查（未连接）"""
        user_id = 999
        
        assert websocket_service.is_user_connected(user_id) is False
    
    def test_get_user_connection_count(self, websocket_service):
        """测试获取特定用户的连接数"""
        user_id = 123
        websocket_service.active_connections[user_id] = [AsyncMock(), AsyncMock()]
        
        if hasattr(websocket_service, 'get_user_connection_count'):
            assert websocket_service.get_user_connection_count(user_id) == 2
        
        # 测试不存在的用户
        if hasattr(websocket_service, 'get_user_connection_count'):
            assert websocket_service.get_user_connection_count(999) == 0
    
    @pytest.mark.asyncio
    async def test_send_market_data_update(self, websocket_service, mock_websocket):
        """测试发送市场数据更新"""
        user_id = 123
        market_data = {
            "symbol": "AAPL",
            "price": Decimal("150.25"),
            "volume": 1000000,
            "timestamp": datetime.now(timezone.utc),
            "change": Decimal("2.50"),
            "change_percent": Decimal("1.69")
        }
        
        # 连接用户
        await websocket_service.connect(mock_websocket, user_id)
        
        # 发送市场数据更新
        if hasattr(websocket_service, 'send_market_data_update'):
            success = await websocket_service.send_market_data_update(user_id, market_data)
            assert success is True
            mock_websocket.send.assert_called_once()
    
    @pytest.mark.asyncio
    async def test_send_subscription_confirmation(self, websocket_service, mock_websocket):
        """测试发送订阅确认"""
        user_id = 123
        symbol = "AAPL"
        
        # 连接用户
        await websocket_service.connect(mock_websocket, user_id)
        
        # 发送订阅确认
        if hasattr(websocket_service, 'send_subscription_confirmation'):
            success = await websocket_service.send_subscription_confirmation(user_id, symbol)
            assert success is True
            mock_websocket.send.assert_called_once()
    
    @pytest.mark.asyncio
    async def test_message_serialization(self, websocket_service, mock_websocket):
        """测试消息序列化"""
        user_id = 123
        
        # 测试包含Decimal和datetime的复杂消息
        complex_message = {
            "type": "market_data",
            "symbol": "AAPL",
            "price": Decimal("150.25"),
            "timestamp": datetime.now(timezone.utc),
            "data": {
                "nested": True,
                "values": [1, 2, 3]
            }
        }
        
        # 连接用户
        await websocket_service.connect(mock_websocket, user_id)
        
        # 发送复杂消息
        success = await websocket_service.send_to_user(user_id, complex_message)
        
        # 验证消息被正确序列化
        assert success is True
        mock_websocket.send.assert_called_once()
        
        # 检查发送的消息是否为有效JSON字符串
        call_args = mock_websocket.send.call_args[0][0]
        assert isinstance(call_args, str)
        
        # 验证可以反序列化
        try:
            parsed_message = json.loads(call_args)
            assert parsed_message["type"] == "market_data"
            assert parsed_message["symbol"] == "AAPL"
        except json.JSONDecodeError:
            pytest.fail("消息序列化后不是有效的JSON")
    
    @pytest.mark.asyncio
    async def test_concurrent_connections(self, websocket_service):
        """测试并发连接处理"""
        # 模拟多个并发连接
        tasks = []
        user_ids = range(1, 11)  # 10个用户
        
        async def connect_user(user_id):
            mock_ws = AsyncMock()
            await websocket_service.connect(mock_ws, user_id)
            return mock_ws
        
        # 并发连接所有用户
        results = await asyncio.gather(*[connect_user(uid) for uid in user_ids])
        
        # 验证所有连接都成功
        assert len(results) == 10
        assert websocket_service.get_active_connections() == 10
        
        # 验证每个用户都有连接
        for user_id in user_ids:
            assert websocket_service.is_user_connected(user_id)
    
    @pytest.mark.asyncio
    async def test_connection_cleanup_on_error(self, websocket_service, mock_websocket):
        """测试错误时的连接清理"""
        user_id = 123
        
        # 连接用户
        await websocket_service.connect(mock_websocket, user_id)
        assert websocket_service.is_user_connected(user_id)
        
        # 模拟WebSocket错误
        mock_websocket.send.side_effect = Exception("WebSocket error")
        
        # 尝试发送消息，应该触发清理
        message = {"type": "test", "data": "test"}
        success = await websocket_service.send_to_user(user_id, message)
        
        # 验证连接被清理
        assert success is False
        # 根据实现，可能需要验证连接是否被移除
    
    @pytest.mark.asyncio
    async def test_memory_leak_prevention(self, websocket_service):
        """测试内存泄漏防护"""
        # 创建大量连接然后断开
        for i in range(100):
            mock_ws = AsyncMock()
            await websocket_service.connect(mock_ws, i)
        
        assert websocket_service.get_active_connections() == 100
        
        # 断开所有连接
        for i in range(100):
            mock_ws = AsyncMock()  # 创建新的mock来模拟断开
            await websocket_service.disconnect(mock_ws, i)
        
        # 验证连接被清理
        # 注意：实际清理可能不是立即的，取决于实现
        remaining_connections = websocket_service.get_active_connections()
        assert remaining_connections <= 100  # 应该显著减少或为0


@pytest.mark.integration
class TestWebSocketServiceIntegration:
    """WebSocket服务集成测试"""
    
    @pytest.mark.asyncio
    async def test_real_websocket_connection(self):
        """测试真实WebSocket连接"""
        # 这个测试需要真实的WebSocket服务器
        # 可以使用websockets库创建测试服务器
        pytest.skip("需要真实WebSocket服务器进行集成测试")
    
    @pytest.mark.asyncio
    async def test_performance_under_load(self, websocket_service):
        """测试负载下的性能"""
        import time
        
        # 测试大量并发连接的性能
        start_time = time.time()
        
        # 创建1000个连接
        tasks = []
        for i in range(1000):
            mock_ws = AsyncMock()
            tasks.append(websocket_service.connect(mock_ws, i))
        
        await asyncio.gather(*tasks)
        
        connection_time = time.time() - start_time
        
        # 验证连接时间在合理范围内（例如，小于5秒）
        assert connection_time < 5.0, f"连接1000个用户耗时过长: {connection_time}秒"
        
        # 测试广播性能
        start_time = time.time()
        message = {"type": "performance_test", "data": "test_data"}
        success_count = await websocket_service.broadcast(message)
        
        broadcast_time = time.time() - start_time
        
        # 验证广播时间和成功率
        assert success_count == 1000
        assert broadcast_time < 2.0, f"广播给1000个用户耗时过长: {broadcast_time}秒"
